'''Define classes and functions for interfacing with SageMaker Ground
Truth object detection.

'''

import os
import imageio
import matplotlib.pyplot as plt
import numpy as np


class BoundingBox:
    '''Bounding box for an object in an image.'''

    def __init__(self, image_id=None, boxdata=None):
        self.image_id = image_id
        if boxdata:
            for datum in boxdata:
                setattr(self, datum, boxdata[datum])

    def __repr__(self):
        return 'Box for image {}'.format(self.image_id)

    def compute_bb_data(self):
        '''Compute the parameters used for IoU.'''
        image = self.image
        self.xmin = self.left/image.width
        self.xmax = (self.left + self.width)/image.width
        self.ymin = self.top/image.height
        self.ymax = (self.top + self.height)/image.height


class WorkerBoundingBox(BoundingBox):
    '''Bounding box for an object in an image produced by a worker.'''

    def __init__(self, image_id=None, worker_id=None, boxdata=None):
        self.worker_id = worker_id
        super().__init__(image_id=image_id, boxdata=boxdata)


class GroundTruthBox(BoundingBox):
    '''Bounding box for an object in an image produced by a worker.'''

    def __init__(self, image_id=None, oiddata=None, image=None):
        self.image = image
        self.class_name = oiddata[0]
        xmin, xmax, ymin, ymax = [float(datum) for datum in oiddata[1:]]
        self.xmin = xmin
        self.ymin = ymin
        self.xmax = xmax
        self.ymax = ymax
        imw = image.width
        imh = image.height
        boxdata = {'height': (ymax-ymin)*imh,
                   'width': (xmax-xmin)*imw,
                   'left': xmin*imw,
                   'top': ymin*imh}
        super().__init__(image_id=image_id, boxdata=boxdata)


class BoxedImage:
    '''Image with bounding boxes.'''

    def __init__(self, id=None, consolidated_boxes=None,
                 worker_boxes=None, gt_boxes=None, uri=None,
                 size=None):
        self.id = id
        self.uri = uri
        if uri:
            self.filename = uri.split('/')[-1]
            self.oid_id = self.filename.split('.')[0]
        else:
            self.filename = None
            self.oid_id = None
        self.local = None
        self.im = None
        if size:
            self.width = size['width']
            self.depth = size['depth']
            self.height = size['height']
            self.shape = self.width, self.height, self.depth
        if consolidated_boxes:
            self.consolidated_boxes = consolidated_boxes
        else:
            self.consolidated_boxes = []
        if worker_boxes:
            self.worker_boxes = worker_boxes
        else:
            self.worker_boxes = []
        if gt_boxes:
            self.gt_boxes = gt_boxes
        else:
            self.gt_boxes = []

    def __repr__(self):
        return 'Image{}'.format(self.id)

    def n_consolidated_boxes(self):
        '''Count the number of consolidated boxes.'''
        return len(self.consolidated_boxes)

    def n_worker_boxes(self):
        return len(self.worker_boxes)

    def download(self, directory):
        target_fname = os.path.join(
            directory, self.uri.split('/')[-1])
        if not os.path.isfile(target_fname):
            os.system(f'aws s3 cp {self.uri} {target_fname}')
        self.local = target_fname

    def imread(self):
        '''Cache the image reading process.'''
        try:
            return imageio.imread(self.local)
        except OSError:
            print("You need to download this image first. "
                  "Use this_image.download(local_directory).")
            raise

    def plot_bbs(self, ax, bbs, img_kwargs, box_kwargs, **kwargs):
        '''Master function for plotting images with bounding boxes.'''
        img = self.imread()
        ax.imshow(img, **img_kwargs)
        imh, imw, *_ = img.shape
        box_kwargs['fill'] = None
        if kwargs.get('worker', False):
            # Give each worker a color.
            worker_colors = {}
            worker_count = 0
            for bb in bbs:
                worker = bb.worker_id
                if worker not in worker_colors:
                    worker_colors[worker] = 'C' + str((9-worker_count) % 10)
                    worker_count += 1
                rec = plt.Rectangle((bb.left, bb.top), bb.width, bb.height,
                                    edgecolor=worker_colors[worker],
                                    **box_kwargs)
                ax.add_patch(rec)
        else:
            for bb in bbs:
                rec = plt.Rectangle(
                    (bb.left, bb.top), bb.width, bb.height, **box_kwargs)
                ax.add_patch(rec)
        ax.axis('off')

    def plot_consolidated_bbs(self, ax, img_kwargs={},
                              box_kwargs={'edgecolor': 'blue',
                                          'lw': 3}):
        '''Plot the consolidated boxes.'''
        self.plot_bbs(ax, self.consolidated_boxes,
                      img_kwargs=img_kwargs, box_kwargs=box_kwargs)

    def plot_worker_bbs(self, ax, img_kwargs={}, box_kwargs={'lw': 2}):
        '''Plot the individual worker boxes.'''
        self.plot_bbs(ax, self.worker_boxes, worker=True,
                      img_kwargs=img_kwargs, box_kwargs=box_kwargs)

    def plot_gt_bbs(self, ax, img_kwargs={},
                    box_kwargs={'edgecolor': 'lime',
                                'lw': 3}):
        '''Plot the ground truth (Open Image Dataset) boxes.'''
        self.plot_bbs(ax, self.gt_boxes,
                      img_kwargs=img_kwargs, box_kwargs=box_kwargs)

    def compute_img_confidence(self):
        ''' Compute the mean bb confidence. '''
        if len(self.consolidated_boxes) > 0:
            return np.mean([box.confidence for box in self.consolidated_boxes])
        else:
            return 0

    def compute_iou_bb(self):
        '''Compute the mean intersection over union for a collection of
        bounding boxes.
        '''

        # Precompute data for the consolidated boxes if necessary.
        for box in self.consolidated_boxes:
            try:
                box.xmin
            except AttributeError:
                box.compute_bb_data()

        # Make the numpy arrays.
        if self.gt_boxes:
            gts = np.vstack([(box.xmin, box.ymin, box.xmax, box.ymax)
                             for box in self.gt_boxes])
        else:
            gts = []
        if self.consolidated_boxes:
            preds = np.vstack([(box.xmin, box.ymin, box.xmax, box.ymax)
                               for box in self.consolidated_boxes])
        else:
            preds = []
        confs = np.array([box.confidence for box in self.consolidated_boxes])

        if len(preds) == 0 and len(gts) == 0:
            return 1.
        if len(preds) == 0 or len(gts) == 0:
            return 0.
        preds = preds[np.argsort(confs.flatten())][::-1]

        is_pred_assigned_to_gt = [False] * len(gts)
        pred_areas = (preds[:, 2] - preds[:, 0]) * \
            (preds[:, 3] - preds[:, 1])
        gt_areas = (gts[:, 2] - gts[:, 0]) * (gts[:, 3] - gts[:, 1])
        all_ious = []
        for pred_id, pred in enumerate(preds):
            best_iou = 0
            best_id = -1
            for gt_id, gt in enumerate(gts):
                if is_pred_assigned_to_gt[gt_id]:
                    continue
                x1 = max(gt[0], pred[0])
                y1 = max(gt[1], pred[1])
                x2 = min(gt[2], pred[2])
                y2 = min(gt[3], pred[3])
                iw = max(0, x2 - x1)
                ih = max(0, y2 - y1)
                inter = iw * ih
                iou = inter / \
                    (pred_areas[pred_id] + gt_areas[gt_id] - inter)
                if iou > best_iou:
                    best_iou = iou
                    best_id = gt_id
            if best_id != -1:
                is_pred_assigned_to_gt[best_id] = True
                # True positive! Store the IoU.
                all_ious.append(best_iou)
            else:
                # 0 IoU for each unmatched gt (false-negative).
                all_ious.append(0.)

        # 0 IoU for each unmatched prediction (false-positive).
        all_ious.extend([0.] * (len(is_pred_assigned_to_gt) -
                                sum(is_pred_assigned_to_gt)))

        return np.mean(all_ious)


def group_miou(imgs):
    '''Compute the mIoU for a group of images.

    Args:
      imgs: list of BoxedImages, with consolidated_boxes and gt_boxes.

    Returns:
      mIoU calculated over the bounding boxes in the group.
    '''
    # Create a notional BoxedImage with bounding boxes from imgs.
    all_consolidated_boxes = [box for img in imgs
                              for box in img.consolidated_boxes]
    all_gt_boxes = [box for img in imgs
                    for box in img.gt_boxes]
    notional_image = BoxedImage(consolidated_boxes=all_consolidated_boxes,
                                gt_boxes=all_gt_boxes)

    # Compute and return the mIoU.
    return notional_image.compute_iou_bb()
